import json
from ai2thor_engine.RocAgent import RocAgent
from ai2thor.controller import Controller
from ai2thor.platform import CloudRendering

from vlm import callVLM
from prompt import *
from utils import *

MAX_MODEL_INFER_COUNT = 10

def main():
    # 启动并复用AI2-THOR控制器
    controller = Controller(
        platform=CloudRendering,
        agent_class="arm",
        snapToGrid=False,
        quality='Medium',
        agentMode="default",
        massThreshold=None,
        scene='FloorPlan1',
        visibilityDistance=20,
        gridSize=0.1,
        renderDepthImage=False,
        renderInstanceSegmentation=False,
        width=800,
        height=450,
        fieldOfView=90,
    )
    # 创建智能体
    autogn = RocAgent(controller, save_path="evaluate/data", scene='FloorPlan1', visibilityDistance=20, gridSize=0.1,
                      fieldOfView=90,
                      target_objects=['CounterTop|-00.08|+01.15|00.00', 'CounterTop|-01.87|+00.95|-01.21'],
                      related_objects=['Apple|-00.47|+01.15|+00.48', 'Pot|-01.22|+00.90|-02.36'],
                      navigable_objects=['Fridge', 'CounterTop', 'CounterTop'],
                      taskid=84,
                      platform_type="GPU")

    # 启动VLM
    vlm_server = callVLM()
    messages = [{
        "role": "user",
        "content": [{"type": "text", "text": EMBODIED_SYSTEM_PROMPT}]
    }]
    print(vlm_server.get_response(messages))

    # 任务目标
    task_name = "get a apple"

    # 任务循环
    call_model_count = 0
    action, raw_action = "init", "init"
    item, pre_item = None, None
    output_response = None
    exec_trajectory = []
    while action != "end" and call_model_count < MAX_MODEL_INFER_COUNT:
        user_txt = ""
        # 重复动作

        # 非法动作
        if check_invalid_action(action, autogn.action_space):
            user_text = INVALID_ACTION_PROMPT.format(action=raw_action)  # action=temp_action
            dic = {
                "response": output_response,
                "action": action,
                "object": item,
                "success": 0,
                "errorInfo": user_text,
                "images": []
            }
            exec_trajectory.append(dic)
            messages.append({
                "role": "user",
                "content": [
                    {"type": "text", "text": user_txt}
                ]
            })
        else:
            print("****** begin exec action:", action, item, "***")
            success, image_fp, legal_locations, legal_objects = autogn.exec(action, item)
            if action == "init":
                success, image_fp, legal_locations, legal_objects = autogn.exec("navigate to", "CounterTop")
            print("****** end exec action:", action, item, "***")
            if not success:
                # exec failed
                print(f"exec <{raw_action}> failed")
                dic = {
                    "response": output_response,            # 上一步模型输出文本
                    "action": action,                       # 执行的动作名
                    "object": item,                         # 动作操作对象
                    "success": 0,                           # 动作执行成功标识符
                    "errorInfo": user_txt,                  # 空字符串，因为执行没有错误
                    "images": image_fp                      # 当前第一视角图像
                }
                exec_trajectory.append(dic)
                messages.append({
                    "role": "user",
                    "content":[
                        {"type": "text", "text": user_txt}
                    ]
                })
                break
            else:
                # exec success
                dic = {
                    "response": output_response,            # 上一步模型输出文本
                    "action": action,                       # 执行的动作名
                    "object": item,                         # 动作操作对象
                    "success": 1,                           # 动作执行成功标识符
                    "errorInfo": "",                        # 空字符串，因为执行没有错误
                    "images": image_fp                      # 执行动作后的图像
                }
                exec_trajectory.append(dic)
                # 基于动作类型生成给模型的消息
                if action == "init":
                    user_txt = TASK_PREFIX_PUT.format(task_name=task_name)
                elif action == action if item is None else action + " " + item:
                    user_txt = USER_IMAGE_PREFIX.format(action=raw_action)
                message = {
                    "role": "user",
                    "content": [
                        {"type": "text", "text": user_txt},
                        {"type": "image", "image": image_fp}
                    ]
                }
                messages.append(message)

        # use VLM
        output_response = vlm_server.get_response(messages)
        call_model_count += 1
        print(output_response)

        # parse response
        raw_action, action, item = macth_action_item(str(output_response), autogn.action_space, {"Apple", "Bread", "Computer"})

        # task is completed
        if raw_action == "end":
            dict = {
                "response": output_response,                    # 上一步模型输出文本
                "action": action,                               # 执行的动作名
                "object": item,                                 # 动作操作对象
                "success": 1,                                   # 动作执行成功标识符
                "errorInfo": "",                                # 空字符串，因为执行没有错误
            }
            exec_trajectory.append(dict)

    # end task exec loop
    if call_model_count == MAX_MODEL_INFER_COUNT:
        dict = {
            "response": output_response,                        # 上一步模型输出文本
            "action": action,                                   # 执行的动作名
            "object": item,                                     # 动作操作对象
            "success": 0,                                       # 动作执行成功标识符
            "errorInfo": "out of MAX_MODEL_INFER_COUNT",        # 执行错误
        }
        exec_trajectory.append(dict)
    # record Trajectory
    with open(f"evaluate/data/result.json","w") as f:
        f.write(json.dumps({
            "trajectory": exec_trajectory,
            "messages": messages
        }, indent=4))

    autogn.controller.stop()
    del autogn

if __name__ == "__main__":
    main()